#include <iostream>
#include <thread>
#include <mutex>
#include <vector>
#include <queue>
#include <random>
#include <chrono>
#include <atomic>
#include <windows.h>

std::mutex cout_mutex;
std::atomic<bool> running(true);
const int WINDOW_SIZE = 4;
const int TOTAL_PACKETS = 15;
const double LOSS_PROBABILITY = 0.2;

struct Packet
{
    int seq_num;
    std::string data;
    bool is_ack;
    bool delivered;
};

struct SenderState
{
    std::vector<Packet> sent_acked;
    std::vector<Packet> sent_unacked;
    std::vector<Packet> can_send;
    std::vector<Packet> cannot_send;
    int next_seq;
    int window_base;
};

struct ReceiverState
{
    std::vector<Packet> acked_delivered;
    std::vector<Packet> can_receive;
    std::vector<Packet> cannot_receive;
    int expected_seq;
};

std::queue<Packet> sender_queue;
std::queue<Packet> receiver_queue;
std::mutex sender_mutex;
std::mutex receiver_mutex;

SenderState sender_state;
ReceiverState receiver_state;
std::mutex state_mutex;

std::random_device rd;
std::mt19937 gen(rd());
std::uniform_real_distribution<> dis(0, 1);

void init_states()
{
    sender_state.next_seq = 0;
    sender_state.window_base = 0;

    for (int i = 0; i < TOTAL_PACKETS; ++i)
    {
        Packet pkt;
        pkt.seq_num = i;
        pkt.data = "Packet " + std::to_string(i);
        pkt.is_ack = false;
        pkt.delivered = false;

        if (i < WINDOW_SIZE)
        {
            sender_state.can_send.push_back(pkt);
        }
        else
        {
            sender_state.cannot_send.push_back(pkt);
        }
    }

    receiver_state.expected_seq = 0;

    for (int i = 0; i < WINDOW_SIZE; ++i)
    {
        Packet pkt;
        pkt.seq_num = -1;
        pkt.is_ack = false;
        receiver_state.can_receive.push_back(pkt);
    }
}

void print_state()
{
    std::lock_guard<std::mutex> lock(cout_mutex);
    system("cls");

    std::cout << "============ Sliding Window Protocol ============" << std::endl;
    std::cout << "Sender State:" << std::endl;
    std::cout << "1. Sent and ACKed: ";
    for (const auto &pkt : sender_state.sent_acked)
    {
        std::cout << pkt.seq_num << " ";
    }
    std::cout << std::endl;

    std::cout << "2. Sent but not ACKed: ";
    for (const auto &pkt : sender_state.sent_unacked)
    {
        std::cout << pkt.seq_num << " ";
    }
    std::cout << std::endl;

    std::cout << "3. Can send: ";
    for (const auto &pkt : sender_state.can_send)
    {
        std::cout << pkt.seq_num << " ";
    }
    std::cout << std::endl;

    std::cout << "4. Cannot send: ";
    for (const auto &pkt : sender_state.cannot_send)
    {
        std::cout << pkt.seq_num << " ";
    }
    std::cout << std::endl;

    std::cout << "\nReceiver State:" << std::endl;
    std::cout << "1. ACKed and delivered: ";
    for (const auto &pkt : receiver_state.acked_delivered)
    {
        std::cout << pkt.seq_num << " ";
    }
    std::cout << std::endl;

    std::cout << "2. Can receive: ";
    for (const auto &pkt : receiver_state.can_receive)
    {
        if (pkt.seq_num != -1)
        {
            std::cout << pkt.seq_num << " ";
        }
        else
        {
            std::cout << "[empty] ";
        }
    }
    std::cout << std::endl;

    std::cout << "3. Cannot receive: ";
    for (const auto &pkt : receiver_state.cannot_receive)
    {
        std::cout << pkt.seq_num << " ";
    }
    std::cout << std::endl;

    std::cout << "\nWindow size: " << WINDOW_SIZE << std::endl;
    std::cout << "Available window: " << (WINDOW_SIZE - sender_state.sent_unacked.size()) << std::endl;
    std::cout << "==========================================" << std::endl;
}

void sender_thread()
{
    while (sender_state.window_base < TOTAL_PACKETS && running)
    {
        std::lock_guard<std::mutex> state_lock(state_mutex);

        while (!sender_state.can_send.empty() && sender_state.sent_unacked.size() < WINDOW_SIZE)
        {
            Packet pkt = sender_state.can_send.front();
            sender_state.can_send.erase(sender_state.can_send.begin());
            sender_state.sent_unacked.push_back(pkt);

            if (dis(gen) >= LOSS_PROBABILITY)
            {
                std::lock_guard<std::mutex> lock(receiver_mutex);
                receiver_queue.push(pkt);

                std::lock_guard<std::mutex> cout_lock(cout_mutex);
                std::cout << "Sender: Sent packet " << pkt.seq_num << std::endl;
            }
            else
            {
                std::lock_guard<std::mutex> cout_lock(cout_mutex);
                std::cout << "Sender: Packet " << pkt.seq_num << " lost" << std::endl;
            }
        }

        {
            std::lock_guard<std::mutex> lock(sender_mutex);
            while (!sender_queue.empty())
            {
                Packet ack = sender_queue.front();
                sender_queue.pop();

                for (auto it = sender_state.sent_unacked.begin(); it != sender_state.sent_unacked.end();)
                {
                    if (it->seq_num == ack.seq_num)
                    {
                        sender_state.sent_acked.push_back(*it);
                        it = sender_state.sent_unacked.erase(it);
                        sender_state.window_base = ack.seq_num + 1;

                        if (!sender_state.cannot_send.empty())
                        {
                            Packet new_pkt = sender_state.cannot_send.front();
                            sender_state.cannot_send.erase(sender_state.cannot_send.begin());
                            sender_state.can_send.push_back(new_pkt);
                        }

                        std::lock_guard<std::mutex> cout_lock(cout_mutex);
                        std::cout << "Sender: Received ACK " << ack.seq_num << std::endl;
                        break;
                    }
                    else
                    {
                        ++it;
                    }
                }
            }
        }

        if (!sender_state.sent_unacked.empty())
        {
            std::this_thread::sleep_for(std::chrono::milliseconds(1500));

            if (!sender_state.sent_unacked.empty())
            {
                std::lock_guard<std::mutex> cout_lock(cout_mutex);
                std::cout << "Sender: Timeout, resending window" << std::endl;

                for (const auto &pkt : sender_state.sent_unacked)
                {
                    if (dis(gen) >= LOSS_PROBABILITY)
                    {
                        std::lock_guard<std::mutex> lock(receiver_mutex);
                        receiver_queue.push(pkt);

                        std::cout << "Sender: Resent packet " << pkt.seq_num << std::endl;
                    }
                    else
                    {
                        std::cout << "Sender: Resent packet " << pkt.seq_num << " lost" << std::endl;
                    }
                }
            }
        }
        else
        {
            std::this_thread::sleep_for(std::chrono::milliseconds(500));
        }
    }

    running = false;
    std::lock_guard<std::mutex> cout_lock(cout_mutex);
    std::cout << "Sender: Transmission complete" << std::endl;
}

void receiver_thread()
{
    while (running)
    {
        std::lock_guard<std::mutex> state_lock(state_mutex);

        {
            std::lock_guard<std::mutex> lock(receiver_mutex);
            while (!receiver_queue.empty())
            {
                Packet pkt = receiver_queue.front();
                receiver_queue.pop();

                if (pkt.seq_num == receiver_state.expected_seq)
                {
                    pkt.delivered = true;
                    receiver_state.acked_delivered.push_back(pkt);
                    receiver_state.expected_seq++;

                    Packet ack;
                    ack.seq_num = pkt.seq_num;
                    ack.is_ack = true;

                    if (dis(gen) >= LOSS_PROBABILITY)
                    {
                        std::lock_guard<std::mutex> sender_lock(sender_mutex);
                        sender_queue.push(ack);

                        std::lock_guard<std::mutex> cout_lock(cout_mutex);
                        std::cout << "Receiver: Received packet " << pkt.seq_num << ", sent ACK" << std::endl;
                    }
                    else
                    {
                        std::lock_guard<std::mutex> cout_lock(cout_mutex);
                        std::cout << "Receiver: ACK " << pkt.seq_num << " lost" << std::endl;
                    }

                    if (!receiver_state.cannot_receive.empty())
                    {
                        Packet new_pkt = receiver_state.cannot_receive.front();
                        receiver_state.cannot_receive.erase(receiver_state.cannot_receive.begin());
                        receiver_state.can_receive.push_back(new_pkt);
                    }
                }
                else
                {
                    std::lock_guard<std::mutex> cout_lock(cout_mutex);
                    std::cout << "Receiver: Out-of-order packet " << pkt.seq_num << ", expected " << receiver_state.expected_seq << ", discarded" << std::endl;
                }
            }
        }

        std::this_thread::sleep_for(std::chrono::milliseconds(300));
    }

    std::lock_guard<std::mutex> cout_lock(cout_mutex);
    std::cout << "Receiver: Exiting" << std::endl;
}

void display_thread()
{
    while (running)
    {
        print_state();
        std::this_thread::sleep_for(std::chrono::seconds(1));
    }
}

int main()
{
    init_states();

    std::thread sender(sender_thread);
    std::thread receiver(receiver_thread);
    std::thread display(display_thread);

    sender.join();
    receiver.join();
    display.join();

    return 0;
}